package com.lm.flight.config;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.apache.arrow.flight.*;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.*;
import org.springframework.context.annotation.Configuration;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;

@Slf4j
@Configuration
public class FlightServerConfig {

    private BufferAllocator allocator;
    private FlightServer flightServer;

    @PostConstruct
    public void startServer() throws Exception {
        allocator = new RootAllocator(Long.MAX_VALUE);

        Location location = Location.forGrpcInsecure("localhost", 8815);
        flightServer = FlightServer.builder(allocator, location, new DemoFlightProducer(allocator)).build();

        flightServer.start();
        log.info("🚀 Arrow Flight Server started on {}", location);
    }

    @PreDestroy
    public void stopServer() throws Exception {
        if (flightServer != null) {
            flightServer.close();
        }
        allocator.close();
        log.info("🛑 Arrow Flight Server stopped.");
    }

    /**
     * 定义一个简单的 FlightProducer，用于返回示例数据
     */
    static class DemoFlightProducer extends NoOpFlightProducer {
        private final BufferAllocator allocator;

        public DemoFlightProducer(BufferAllocator allocator) {
            this.allocator = allocator;
        }

        @Override
        public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
            // 定义字段（id:int, name:varchar）
            Field idField = new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null);
            Field nameField = new Field("name", FieldType.nullable(new ArrowType.Utf8()), null);
            List<Field> fields = Arrays.asList(idField, nameField);
            Schema schema = new Schema(fields);

            // 创建 VectorSchemaRoot
            try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
                listener.start(root);
                root.allocateNew();

                // 获取具体类型向量并分配内存
                IntVector idVector = (IntVector) root.getVector("id");
                VarCharVector nameVector = (VarCharVector) root.getVector("name");

                // 写入数据（示例两行）
                idVector.setSafe(0, 1);
                nameVector.setSafe(0, "Alice".getBytes(StandardCharsets.UTF_8));

                idVector.setSafe(1, 2);
                nameVector.setSafe(1, "Bob".getBytes(StandardCharsets.UTF_8));

                // 设置有效行数（两种方式均可）
                root.setRowCount(2);
                // 或者单独设置每个向量的 value count：
                // idVector.setValueCount(2);
                // nameVector.setValueCount(2);

                // 发送数据
                listener.putNext();
                listener.completed();
            } catch (Exception e) {
                listener.error(e);
            }
        }
    }
}